import numpy as np
import matplotlib.pyplot as plt

# 参数设置
T_load = 0.03  # 荷载作用时长 (秒)
peak = 1.0
omega0 = np.pi / T_load  # 角频率 (rad/s)

T_total = 21 * T_load    # 总时长
N = 256                 # 采样点数
dt = T_total / N        # 采样间隔 (秒)
t = np.linspace(0, T_total, N, endpoint=False)  # 时间序列

# 定义荷载作用区间：10T ≤ t ≤ 11T
St = 10 * T_load  # 起始时间
mask = (t >= St) & (t <= St + T_load)
s = np.where(mask, peak * np.sin(np.pi * (t - St) / T_load), 0.0)

# 数值傅里叶变换
fft_vals = np.fft.fft(s) * dt  # 乘以dt近似连续积分
freq = np.fft.fftfreq(N, dt)
omega = 2 * np.pi * freq  # 转换为角频率
fft_shifted = np.fft.fftshift(fft_vals)
freq_shifted = np.fft.fftshift(freq)
omega_shifted = 2 * np.pi * freq_shifted  # 转换为角频率

# 解析傅里叶变换计算
omega = omega_shifted.copy()
T0 = T_load
St = 10 * T0

# 分子和分母计算
numerator = peak * np.pi * T0 * (np.exp(-1j * omega * T0) + 1) * np.exp(-1j * omega * St)
denominator = np.pi**2 - (omega * T0)**2
singular_mask = np.isclose(np.abs(denominator), 0, atol=1e-10)
S_analytical = np.divide(numerator, denominator, where=~singular_mask)

# 精确处理每个奇异点（修正赋值逻辑）
mask_plus_pi = np.isclose(omega * T0, np.pi, atol=1e-8)
mask_minus_pi = np.isclose(omega * T0, -np.pi, atol=1e-8)

if np.any(mask_plus_pi):
    S_analytical[mask_plus_pi] = (-1j * peak * T0 / 2) * np.exp(-1j * omega[mask_plus_pi] * St)

if np.any(mask_minus_pi): 
    S_analytical[mask_minus_pi] = (-1j * peak * T0 / 2) * np.exp(-1j * omega[mask_minus_pi] * St)

# 绘图
plt.figure(figsize=(10, 6), dpi=300)
'''
# 时域信号
plt.subplot(2, 2, 1)
plt.plot(t, s)
plt.xlabel('Time (s)')
plt.ylabel('Amplitude')
plt.title('Time Domain Signal')
plt.xlim(St - T0, St + 2*T0)  # 聚焦荷载作用区间
'''
# 频域幅度（数值与解析对比）
plt.subplot(2, 2, 2)
plt.plot(freq_shifted, np.abs(S_analytical), 'b', lw=2, label='Analytical')
plt.plot(freq_shifted, np.abs(fft_shifted), 'r--', label='Numerical')
#plt.plot(freq, np.abs(fft_vals), label='Numerical')
#plt.plot(freq, np.abs(S_analytical), 'r--', lw=2, label='Analytical')
plt.xlabel('Frequency (Hz)', fontsize=10)
plt.ylabel(r"|$\tilde{p}_t$|", fontsize=10)
plt.legend()
plt.title('Frequency Domain Comparison', fontsize=12)
'''
# 频域相位（数值与解析对比）
plt.subplot(2, 2, 3)
plt.plot(freq_shifted, np.angle(fft_shifted), '.', label='Numerical')
plt.plot(freq_shifted, np.angle(S_analytical), 'r-', lw=2, label='Analytical')
#plt.plot(freq, np.angle(fft_vals), '.', label='Numerical')
#plt.plot(freq, np.angle(S_analytical), 'r-', lw=2, label='Analytical')
plt.xlabel('Frequency (Hz)')
plt.ylabel('Phase (rad)')
plt.legend()
plt.title('Frequency Domain (Phase)')
'''
# 逆变换对比（数值方法）
#s_reconstructed = np.fft.ifft(np.fft.ifftshift(fft_shifted / dt)).real
s_reconstructed = np.fft.ifft(fft_vals / dt).real
plt.subplot(2, 2, 1)
plt.plot(t, s, 'b', label='Original')
plt.plot(t, s_reconstructed, 'r--', label='Reconstructed')
plt.xlabel('Time (s)', fontsize=10)
plt.ylabel('$p_t$ (Pa)', fontsize=10)
plt.legend(loc='upper right')
plt.title('Time Domain Comparison', fontsize=12)
plt.xlim(St - T0, St + 2*T0)

plt.tight_layout()
plt.savefig('FFT.png', bbox_inches='tight')
plt.show()